/**
 * Copyright 2008 Adam Ruggles.
 * 
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 * http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package kiff;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.sql.SQLException;

import javax.sql.DataSource;

import org.dbunit.DatabaseUnitException;
import org.dbunit.database.DatabaseConfig;
import org.dbunit.database.DatabaseConnection;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.xml.FlatDtdDataSet;
import org.dbunit.dataset.xml.FlatXmlDataSet;
import org.dbunit.ext.hsqldb.HsqldbDataTypeFactory;
import org.dbunit.operation.DatabaseOperation;
import org.springframework.jdbc.datasource.DataSourceUtils;

/**
 * Data Utilities Class.  Used to push testing data into a data source for running tests.
 * @author Adam
 */
public final class DataUtil {
	/**
	 * XML Pointing to the data to push to the database before the tests.
	 */
	private static final String DEFAULT_DATA_DTD_FILE = "src/test/resources/dataset.dtd";

	/**
	 * XML Pointing to the data to push to the database before the tests.
	 */
	private static final String DEFAULT_DATA_XML_FILE = "src/test/resources/dataset.xml";

	/**
	 * Creates the DTD file for the given datasource.
	 * @param dataSource The DataSource to use to connect to the database.
	 * @throws IOException If an I/O error occurs creating the DTD.
	 * @throws SQLException If a SQL error occurs creating the DTD.
	 * @throws DatabaseUnitException If a DB Unit error occurs creating the DTD.
	 */
	public static void createDTD(final DataSource dataSource) throws IOException, SQLException, DatabaseUnitException {
		// Get the Database Connection.
		IDatabaseConnection connection = new DatabaseConnection(DataSourceUtils.getConnection(dataSource));
		// Set the Data Type Factory.
		connection.getConfig().setProperty(DatabaseConfig.PROPERTY_DATATYPE_FACTORY, new HsqldbDataTypeFactory());
		// write DTD to the dtdStream
		ByteArrayOutputStream dtdStream = new ByteArrayOutputStream();
        FlatDtdDataSet.write(connection.createDataSet(), dtdStream);
		// Write the DTD to a file.
        File file = new File(DEFAULT_DATA_DTD_FILE);
        FileOutputStream fos = new FileOutputStream(file); 
        dtdStream.writeTo(fos); 
        fos.flush();
        fos.close(); 
	}

	/**
	 * Adds the test data using the default dataset.xml file to the database.
	 * @param dataSource The DataSource to use to connect to the database.
	 * @throws IOException If an I/O error occurs adding the test data.
	 * @throws SQLException If a SQL error occurs adding the test data.
	 * @throws DatabaseUnitException If a DB Unit error occurs adding the test data.
	 */
	public static void createTestData(final DataSource dataSource)
		throws IOException, SQLException, DatabaseUnitException {
		createTestData(dataSource, getResourceAsStream(DEFAULT_DATA_XML_FILE));
	}

	/**
	 * Adds the test data in the specified filename to the database.
	 * @param dataSource The DataSource to use to connect to the database.
	 * @param file The XML file containing the data set.
	 * @throws IOException If an I/O error occurs adding the test data.
	 * @throws SQLException If a SQL error occurs adding the test data.
	 * @throws DatabaseUnitException If a DB Unit error occurs adding the test data.
	 * @throws FileNotFoundException If the file does not exist.
	 */
	public static void createTestData(final DataSource dataSource, final File file)
		throws FileNotFoundException, IOException, SQLException, DatabaseUnitException {
        createTestData(dataSource, new FileInputStream(file));
	}

	/**
	 * Adds the test data using a data set input stream.
	 * @param dataSource The DataSource to use to connect to the database.
	 * @param inputStream The data set input stream.
	 * @throws IOException If an I/O error occurs adding the test data.
	 * @throws SQLException If a SQL error occurs adding the test data.
	 * @throws DatabaseUnitException If a DB Unit error occurs adding the test data.
	 */
	public static void createTestData(final DataSource dataSource, final InputStream inputStream)
		throws IOException, SQLException, DatabaseUnitException {
		// Get the Database Connection.
		IDatabaseConnection connection = new DatabaseConnection(DataSourceUtils.getConnection(dataSource));
		// Set the Data Type Factory.
		connection.getConfig().setProperty(DatabaseConfig.PROPERTY_DATATYPE_FACTORY, new HsqldbDataTypeFactory());
		// write DTD to the dtdStream
		ByteArrayOutputStream dtdStream = new ByteArrayOutputStream();
        FlatDtdDataSet.write(connection.createDataSet(), dtdStream);
       
        ByteArrayInputStream dtdInputStream = new ByteArrayInputStream(dtdStream.toByteArray());
		IDataSet dataSet = new FlatXmlDataSet(inputStream, dtdInputStream);
		try {
			dataSource.getConnection().createStatement().execute("SET REFERENTIAL_INTEGRITY FALSE;");
            DatabaseOperation.CLEAN_INSERT.execute(connection, dataSet);
            dataSource.getConnection().createStatement().execute("SET REFERENTIAL_INTEGRITY TRUE;");
        } finally {
            connection.close();
        }
	}

	/**
	 * Adds the test data in the default data.xml file to the database.
	 * @param dataSource The DataSource to use to connect to the database.
	 * @param filename The path to the XML file containing the data set.
	 * @throws Exception If an error occurs adding the test data.
	 */
	public static void createTestData(final DataSource dataSource, final String filename) throws Exception {
		createTestData(dataSource, getResourceAsStream(filename));
	}

	/**
	 * Removes all the test data from the database.
	 * @param dataSource The DataSource to use to connect to the database.
	 * @throws IOException If an I/O error occurs adding the test data.
	 * @throws SQLException If a SQL error occurs adding the test data.
	 * @throws DatabaseUnitException If a DB Unit error occurs adding the test data.
	 */
	public static void deleteTestData(final DataSource dataSource) 
		throws IOException, SQLException, DatabaseUnitException {
		// Get the Database Connection.
		IDatabaseConnection connection = new DatabaseConnection(DataSourceUtils.getConnection(dataSource));
		// Set the Data Type Factory.
		connection.getConfig().setProperty(DatabaseConfig.PROPERTY_DATATYPE_FACTORY, new HsqldbDataTypeFactory());
		// write DTD to the dtdStream
		ByteArrayOutputStream dtdStream = new ByteArrayOutputStream();
        FlatDtdDataSet.write(connection.createDataSet(), dtdStream);
       
        ByteArrayInputStream dtdInputStream = new ByteArrayInputStream(dtdStream.toByteArray());
		IDataSet dataSet = new FlatXmlDataSet(getResourceAsStream(DEFAULT_DATA_XML_FILE), dtdInputStream);
		try {
			dataSource.getConnection().createStatement().execute("SET REFERENTIAL_INTEGRITY FALSE;");
            DatabaseOperation.DELETE_ALL.execute(connection, dataSet);
            dataSource.getConnection().createStatement().execute("SET REFERENTIAL_INTEGRITY TRUE;");
        } finally {
            connection.close();
        }
	}

	/**
	 * Returns a file resource as an input stream.
	 * @param fileName The name of the resource or path to the file.
	 * @return A file resource as an input stream.
	 * @throws IllegalArgumentException If the file resource could not be found.
	 */
	public static InputStream getResourceAsStream(final String fileName) throws IllegalArgumentException {
		InputStream stream = Thread.currentThread().getContextClassLoader().getResourceAsStream(fileName);
		if (stream == null) {
			stream = ClassLoader.getSystemResourceAsStream(fileName);
		}
		if (stream == null) {
			File file = new File(fileName);
			try	{
				stream = new FileInputStream(file);
			} catch (FileNotFoundException fileNotFoundEx) {
				throw new IllegalArgumentException("File not found.", fileNotFoundEx);
			}
		}
		return stream;
	}

	/**
	 * Private Constructor.
	 */
	private DataUtil() { }
}
